package com.zhang.sparksql_1

import org.apache.spark.SparkConf
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, LongType, StructField, StructType}
import org.apache.spark.sql.{DataFrame, Row, SparkSession}

/**
 * @title:
 * @author: zhang
 * @date: 2021/12/10 19:48 
 */
object SparkSQl_UDAF_03 {
  def main(args: Array[String]): Unit = {
    //  获取SparkSession连接对象
    val conf = new SparkConf().setMaster("local[*]").setAppName("spark-sql")
    val spark: SparkSession = SparkSession.builder().config(conf).getOrCreate()
    //todo 创建DataFrame
    val df: DataFrame = spark.read.json("datas/user.json")

    df.createOrReplaceTempView("user")

    //todo 创建UDAF函数
    spark.udf.register("MyAvg",new MyAvgUDAF())
    spark.sql("select MyAvg(age) from user").show()
    //todo 关闭资源
    spark.stop()
  }

  class MyAvgUDAF extends UserDefinedAggregateFunction{
    override def inputSchema: StructType = {
      StructType(
        Array(
          StructField("age",LongType)
        )
      )
    }

    override def bufferSchema: StructType = {
      StructType(
        Array(
          StructField("total",LongType),
          StructField("count",LongType)
        )
      )
    }

    override def dataType: DataType = LongType

    override def deterministic: Boolean = true

    override def initialize(buffer: MutableAggregationBuffer): Unit = {
      buffer.update(0,0L)
      buffer.update(1,0L)
    }

    override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
      buffer.update(0,buffer.getLong(0)+input.getLong(0))
      buffer.update(1,buffer.getLong(1)+1)
    }

    override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
      buffer1.update(0,buffer1.getLong(0)+buffer2.getLong(0))
      buffer1.update(1,buffer1.getLong(1)+buffer2.getLong(1))
    }

    override def evaluate(buffer: Row): Any = {
      buffer.getLong(0)/buffer.getLong(1)
    }
  }


}
